from pandas import DataFrameMapper
from sklearn.preprocessing import OneHotEncoder
from sklearn2pmml.decoration import ContinuousDomain, CategoricalDomain
from sklearn.pipeline import Pipeline

import shap

def onehot_pipeline(model, X_train, y_train, char_cols=None, num_fillna=None, char_fillna=None):
    '''
    传入带有参数的模型,封装成类别特征one-hot的pipline
    ————————————————————————————————————
    入参:
        model:带有参数的模型
        X_train:训练集的特征,pd.DataFrame格式
        Y_train:训练集的目标
        char_cols:类别特征的列表,如不传入自动根据数据类型获取
        num_fillna:数值特征的缺失填充值,可支持不填充
        char_fillna:类别特征的缺失填充值,可支持不传入,但模型会自动填充null用于one-hot
    出参:
        pipeline:封装好mapper和model的pipeline,并训练完成
    '''
    if not char_cols:
        col_types = X_train.dtypes
        char_cols = list(col_types[col_types.apply(lambda x: 'int' not in str(x) and 'float' not in str(x))].index)
    num_cols = list(set(X_train.columns) - set(char_cols))
    
    if not isinstance(char_fillna, str):
        char_fillna = 'null'
        
    mapper = DataFrameMapper(
        [(num_cols, ContinuousDomain(missing_value_replacement=num_fillna, with_data=False))] +
        [([char_col], [CategoricalDomain(missing_value_replacement=char_fillna, invalid_value_treatment='as_is'),
                       OneHotEncoder(handle_unknown='ignore')])
         for char_col in char_cols]
    )
    
    pipeline = Pipeline(steps=[('mapper', mapper), ('model', model)])
    pipeline.fit(X_train, y_train)
    
    return pipeline

def pipeline_shap(pipeline, X_train, y_train, interaction=False, sample=None):
    '''
    获取由onehot_pipeline返回的pipeline的shap值
    ————————————————————————————————————
    入参:
        pipeline:onehot_pipeline的返回对象
        X_train:训练集的特征,pd.DataFrame格式
        Y_train:训练集的目标
        interaction:是否返回shap interaction values
        sample:抽样数int或抽样比例float,不传入则不抽样
    出参:
        feature_values:如传入sample则是抽样后的X_train,否则为X_train
        shap_values:pd.DataFrame格式shap values,如interaction传入True,则为shap interaction values
    '''
    
    
    if isinstance(sample, int):
        feature_values = X_train.sample(n=sample)
    elif isinstance(sample, float):
        feature_values = X_train.sample(frac=sample)
    else:
        feature_values = X_train
        
    mapper = pipeline.steps[0][1]
    model = pipeline._final_estimator
    sort_cols, onehot_cols = [], []
    for i in mapper.features:
        sort_cols += i[0]
        if 'OneHot' in str(i[1]):
            onehot_cols += i[0]
    feature_values = feature_values[sort_cols]
    
    mapper.fit(X_train)
    X_train_mapper = mapper.transform(X_train)
    feature_values_mapper = mapper.transform(feature_values)
    model.fit(X_train_mapper, y_train)
    
    shap_values = pd.DataFrame(index=feature_values.index, columns=feature_values.columns)
    explainer = shap.TreeExplainer(model)
    if interaction:
        mapper_shap_values = explainer.shap_interaction_values(feature_values_mapper)
        col_index = 0
        for col in sort_cols:
            if col in onehot_cols:
                col_index_span = len(X_train[col].unique())
                shap_values[col] = mapper_shap_values[
                    :, col_index: col_index + col_index_span, col_index: col_index + col_index_span
                ].sum(2).sum(1)
                col_index += col_index_span
            else:
                shap_values[col] = mapper_shap_values[:, col_index, col_index]
                col_index += 1
    else:
        mapper_shap_values = explainer.shap_values(feature_values_mapper)
        if len(mapper_shap_values) == 2:
            mapper_shap_values = mapper_shap_values[1]
        col_index = 0
        for col in sort_cols:
            if col in onehot_cols:
                col_index_span = len(X_train[col].unique())
                shap_values[col] = mapper_shap_values[
                    :, col_index: col_index + col_index_span
                ].sum(1)
                col_index += col_index_span
            else:
                shap_values[col] = mapper_shap_values[:, col_index]
                col_index += 1
                
    return feature_values, shap_values